import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.ticker import MaxNLocator
from matplotlib.colors import ListedColormap, LinearSegmentedColormap

def plotSpatial(ax, XTicks, YTicks, ResultList):
   #set colormap
    colors2 = ["#fbfcf7","#6aa03f","#0e6bb8","#7c2e97","#fbfcf7"]
    cmap2 = LinearSegmentedColormap.from_list("mycmap", colors2)
    nx = len(XTicks)
    ny = len(YTicks)

    #plot
    im = ax.imshow(ResultList, aspect = "auto", cmap=cmap2)

    # 添加颜色条
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="2%", pad=0.1)  # 设置 colorbar 的宽度和位置
    cb = plt.colorbar(im, cax=cax)
    # cb = plt.colorbar(im, cax=cax, label='Color Scale')

    # 设置坐标轴
    ax.set_xticks(np.arange(0, nx, 1), XTicks)
    ax.set_yticks(np.arange(0, ny, 1), YTicks)
    ax.xaxis.set_major_locator(MaxNLocator(nbins=10))
    ax.yaxis.set_major_locator(MaxNLocator(nbins=10))

    # 设置其他标签和标题等
    ax.set_xlabel('distance [$\mu$m]', fontsize=10)
    ax.set_ylabel('radius [$\mu$m]', fontsize=10)
    ax.set_title("phase", fontsize=10)



def plotAxial(ax, XTicks, ResultList):
    #plot
    ax.plot(XTicks, ResultList)  
    # 设置其他标签和标题等
    ax.set_xlabel('distance [$\mu$m]', fontsize=10)
    ax.set_ylabel('phase', fontsize=10)
    ax.set_title("phase", fontsize=10)
